import json

from dotenv import load_dotenv
from vanna.telchina.telchina_chat import Telchina_SQL_LocalContext_OpenAI, Telchina_Command_LocalContext_OpenAI

from request_util import RequestUtil

load_dotenv()

from functools import wraps
from flask import Flask, jsonify, Response, request, redirect, url_for
import flask
import os
from cache import MemoryCache

app = Flask(__name__, static_url_path='', static_folder='../../static')

# SETUP
cache = MemoryCache()
from openai import OpenAI
client = OpenAI(
    api_key=os.environ['API_KEY'],
    base_url=os.environ['BASE_URL']
)
config={
    "model": os.environ['MODEL'],
    "temperature": os.environ['TEMPERATURE'],
    "max_tokens": os.environ['MAX_TOKENS'],
    #"initial_prompt": "你是一个数据库专家"
}
vn = Telchina_SQL_LocalContext_OpenAI(client=client, config=config)
vn.connect_to_postgres()

vn_command = Telchina_Command_LocalContext_OpenAI(client=client, config=config)

# NO NEED TO CHANGE ANYTHING BELOW THIS LINE
def requires_cache(fields):
    def decorator(f):
        @wraps(f)
        def decorated(*args, **kwargs):
            id = request.args.get('id')

            if id is None:
                return jsonify({"type": "error", "error": "No id provided"})

            for field in fields:
                if cache.get(id=id, field=field) is None:
                    return jsonify({"type": "error", "error": f"No {field} found"})

            field_values = {field: cache.get(id=id, field=field) for field in fields}

            # Add the id to the field_values
            field_values['id'] = id

            return f(*args, **field_values, **kwargs)
        return decorated
    return decorator

@app.route('/api/v0/generate_questions', methods=['GET'])
def generate_questions():
    return jsonify({
        "type": "question_list",
        "questions": vn.generate_questions(),
        "header": "以下是您可以问的一些问题："
        })

@app.route('/api/v0/generate_sql', methods=['GET'])
def generate_sql():
    question = flask.request.args.get('question')

    if question is None:
        return jsonify({"type": "error", "error": "没有提供问题"})

    id = cache.generate_id(question=question)
    sql = vn.generate_sql(question=question)

    cache.set(id=id, field='question', value=question)
    cache.set(id=id, field='sql', value=sql)

    return jsonify(
        {
            "type": "sql",
            "id": id,
            "text": sql,
        })

@app.route('/api/v0/run_sql', methods=['GET'])
@requires_cache(['sql'])
def run_sql(id: str, sql: str):
    try:
        vn.connect_to_postgres()

        df = vn.run_sql(sql=sql)

        cache.set(id=id, field='df', value=df)

        return jsonify(
            {
                "type": "df",
                "id": id,
                #"df": df.head(10).to_json(orient='records'),
                "df": df.to_json(orient='records'),
            })

    except Exception as e:
        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/download_csv', methods=['GET'])
@requires_cache(['df'])
def download_csv(id: str, df):
    csv = df.to_csv()

    return Response(
        csv,
        mimetype="text/csv",
        headers={"Content-disposition":
                 f"attachment; filename={id}.csv"})

@app.route('/api/v0/generate_plotly_figure', methods=['GET'])
@requires_cache(['df', 'question', 'sql'])
def generate_plotly_figure(id: str, df, question, sql):
    try:
        code = vn.generate_plotly_code(question=question, sql=sql, df_metadata=f"Running df.dtypes gives:\n {df.dtypes}")
        fig = vn.get_plotly_figure(plotly_code=code, df=df, dark_mode=False)
        fig_json = fig.to_json()

        cache.set(id=id, field='fig_json', value=fig_json)

        return jsonify(
            {
                "type": "plotly_figure",
                "id": id,
                "fig": fig_json,
            })
    except Exception as e:
        # Print the stack trace
        import traceback
        traceback.print_exc()

        return jsonify({"type": "error", "error": str(e)})


# @app.route('/api/v0/generate_plotly_figure', methods=['GET'])
# @requires_cache(['df', 'question', 'sql'])
# def generate_plotly_figure(id: str, df, question, sql):
#     try:
#         # code = vn.generate_plotly_code(question=question, sql=sql, df_metadata=f"Running df.dtypes gives:\n {df.dtypes}")
#         # fig = vn.get_plotly_figure(plotly_code=code, df=df, dark_mode=False)
#         # fig_json = fig.to_json()
#         #
#         # cache.set(id=id, field='fig_json', value=fig_json)
#
#         return jsonify(
#             {
#                 "type": "plotly_figure",
#                 "id": id,
#                 "fig": '{}',
#             })
#     except Exception as e:
#         # Print the stack trace
#         import traceback
#         traceback.print_exc()
#
#         return jsonify({"type": "error", "error": str(e)})



@app.route('/api/v0/get_training_data', methods=['GET'])
def get_training_data():
    df = vn.get_training_data()

    return jsonify(
    {
        "type": "df",
        "id": "training_data",
        "df": df.head(25).to_json(orient='records'),
    })

@app.route('/api/v0/remove_training_data', methods=['POST'])
def remove_training_data():
    # Get id from the JSON body
    id = flask.request.json.get('id')

    if id is None:
        return jsonify({"type": "error", "error": "No id provided"})

    if vn.remove_training_data(id=id):
        return jsonify({"success": True})
    else:
        return jsonify({"type": "error", "error": "Couldn't remove training data"})

@app.route('/api/v0/train', methods=['POST'])
def add_training_data():
    question = flask.request.json.get('question')
    sql = flask.request.json.get('sql')
    ddl = flask.request.json.get('ddl')
    documentation = flask.request.json.get('documentation')

    try:
        id = vn.train(question=question, sql=sql, ddl=ddl, documentation=documentation)

        return jsonify({"id": id})
    except Exception as e:
        print("TRAINING ERROR", e)
        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/generate_followup_questions', methods=['GET'])
@requires_cache(['df', 'question', 'sql'])
def generate_followup_questions(id: str, df, question, sql):
    #followup_questions = vn.generate_followup_questions(question=question, sql=sql, df=df)

    #cache.set(id=id, field='followup_questions', value=followup_questions)

    return jsonify(
        {
            "type": "question_list",
            "id": id,
            #"questions": followup_questions,
            "questions": '{}',
            "header": "Here are some followup questions you can ask:"
        })

# @app.route('/api/v0/load_question', methods=['GET'])
# @requires_cache(['question', 'sql', 'df', 'fig_json', 'followup_questions'])
# def load_question(id: str, question, sql, df, fig_json, followup_questions):
#     try:
#         return jsonify(
#             {
#                 "type": "question_cache",
#                 "id": id,
#                 "question": question,
#                 "sql": sql,
#                 "df": df.head(10).to_json(orient='records'),
#                 "fig": fig_json,
#                 "followup_questions": followup_questions,
#             })
#
#     except Exception as e:
#         return jsonify({"type": "error", "error": str(e)})



@app.route('/api/v0/load_question', methods=['GET'])
@requires_cache(['question', 'sql', 'df'])
def load_question(id: str, question, sql, df):
    try:
        return jsonify(
            {
                "type": "question_cache",
                "id": id,
                "question": question,
                "sql": sql,
                "df": df.head(10).to_json(orient='records'),
                "fig": "{}",
                "followup_questions": {}
            })

    except Exception as e:
        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/get_question_history', methods=['GET'])
def get_question_history():
    return jsonify({"type": "question_history", "questions": cache.get_all(field_list=['question']) })




@app.route('/api/v0/generate_command', methods=['GET'])
def generate_command():
    """
    {"type": "command", "error": "没有提供问题"}
    {"type": "command", "message": "指令操作成功"}
    """
    question = flask.request.args.get('question')
    if question is None:
        return jsonify({"type": "command", "error": "没有提供问题"})
    res = vn_command.generate_command(question)
    commandObj = json.loads(res)
    if commandObj.get("error") is not None:
        return jsonify({"type": "command", "error": commandObj["error"]})
    command = commandObj["command"]
    lampPostCode = commandObj["lampPostCode"]
    dimmingValue = commandObj["dimmingValue"]
    on_off = '开灯' if command == 'on' else '关灯'
    #根据灯杆查询
    sql = f'''
    SELECT
        t2.alone_lamp_control_id,
        t3.alone_lamp_id,
        t4.alone_lamp_code,
        {dimmingValue} AS dimming_value
    FROM
        device.l_lamppost t1
        JOIN device.l_alonelamp_control t2 ON t2.lamp_post_id = t1.lamp_post_id
        JOIN device.l_alonelamp_control_middle t3 ON t3.alone_lamp_control_id = t2.alone_lamp_control_id
        JOIN device.l_alonelamp t4 ON t4.alone_lamp_id = t3.alone_lamp_id
    WHERE
        t1.lamp_post_code = '{lampPostCode}'
    '''
    #print(f'请求查询sql:{sql}')
    vn.connect_to_postgres()
    data = vn.run_sql(sql=sql)
    if data is not None and data.size > 0:
        try:
            params = []
            for index, row in data.iterrows():
                params.append({
                    "aloneLampControlId": row["alone_lamp_control_id"],
                    "aloneLampId": row["alone_lamp_id"],
                    "aloneLampCode": row["alone_lamp_code"],
                    "dimmingValue": row["dimming_value"]
                })
            # 发送调光指令
            util = RequestUtil()
            token = 'eyJhbGciOiJIUzI1NiJ9.eyJqdGkiOiJjb20udGVsY2hpbmEuand0Y2VudGVyIiwiaWF0IjoxNjAxMjgxODY1LCJzdWIiOiJhZG1pbiIsImlzcyI6ImNvbS50ZWxjaGluYS5qd3RjZW50ZXIifQ.oY4eoFe23dSpXnfObLJrovu0fC_ZkwScQwjbkE82Ow8'
            response_post = util.post("http://123.232.113.233:55100/smartpole/light/command/setDimmingValue", json=params, headers={"X-Authentication-Token": token})
            if response_post is not None :
                message = f"{lampPostCode}{on_off}指令发送成功"
                return jsonify({"type": "command", "message": message})
        except Exception as e:
            print(f"请求失败: {e}")
            message = f"{lampPostCode}{on_off}指令发送失败"
            return jsonify({"type": "command", "message": message})
    else:
        message = f"未找到灯杆:{lampPostCode}，{on_off}指令发送失败"
        return jsonify({"type": "command", "message": message})

@app.route('/')
def root():
    return app.send_static_file('index.html')

if __name__ == '__main__':
    app.run(debug=True,host='0.0.0.0')
